import os
import numpy as np
import torch
import matplotlib.pyplot as plt
import random
import time
from tqdm import tqdm
import argparse
import sys

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from env.grid_env import setup_chinese_font, GridMap, RobotEnv
from agent.dqn_agent import DQNAgent

# 设置中文字体
setup_chinese_font()

def set_seed(seed):
    """设置随机种子"""
    if seed is not None:
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)

def create_test_map(width=8, height=8, seed=None):
    """创建一个测试用地图，包含一些特殊的障碍物布局，容易导致振荡行为"""
    grid_map = GridMap(width=width, height=height, obstacle_ratio=0, seed=seed)
    
    # 手动设置一些障碍物，创造容易产生振荡的环境
    obstacles = []
    
    # 添加一些走廊和狭窄通道
    for i in range(2, 6):
        obstacles.append((i, 2))
        obstacles.append((i, 5))
    
    for i in range(3, 5):
        obstacles.append((2, i))
        obstacles.append((5, i))
    
    # 添加一些角落和死胡同
    obstacles.append((1, 1))
    obstacles.append((width-2, height-2))
    obstacles.append((width-2, 1))
    obstacles.append((1, height-2))
    
    # 设置障碍物
    grid_map.set_obstacles(obstacles)
    return grid_map

def train_oscillation_fix(episodes=100, use_fix=True, seed=42):
    """
    训练DQN智能体，测试振荡修复效果
    
    Args:
        episodes: 训练回合数
        use_fix: 是否使用振荡修复
        seed: 随机种子
    """
    # 设置随机种子
    set_seed(seed)
    
    # 创建保存目录
    os.makedirs("results", exist_ok=True)
    
    # 创建测试地图
    grid_map = create_test_map(seed=seed)
    env = RobotEnv(grid_map)
    
    # 创建DQN智能体
    state_dim = 15 if use_fix else 10  # 使用修复时包含振荡特征和历史动作
    action_dim = 4  # [上, 右, 下, 左]
    
    # DQN配置
    dqn_config = {
        'learning_rate': 0.001,
        'gamma': 0.99,
        'epsilon_start': 1.0,
        'epsilon_end': 0.05,
        'epsilon_decay': 0.95,
        'batch_size': 32,
        'memory_size': 10000,
        'target_update': 10
    }
    
    # 创建智能体
    agent = DQNAgent(
        state_dim=state_dim,
        action_dim=action_dim,
        config=dqn_config,
        use_double_dqn=True,
        use_prioritized=True
    )
    
    # 记录训练数据
    rewards_history = []
    steps_history = []
    success_history = []
    oscillation_counts = []  # 记录每回合振荡次数
    
    # 设置动态可视化
    plt.ion()  # 打开交互模式
    fig, ax = plt.subplots(figsize=(10, 10))
    plt.show()
    
    # 开始训练
    print(f"开始训练DQN智能体，{'使用' if use_fix else '不使用'}振荡修复，共{episodes}回合...")
    
    for episode in range(1, episodes + 1):
        # 确保起点和目标点之间有可达路径
        state = None
        while state is None:
            try:
                state = env.reset(ensure_path=True)
            except ValueError:
                continue
        
        episode_reward = 0
        num_steps = 0
        oscillation_count = 0  # 当前回合振荡次数
        
        # 记录路径用于可视化
        path = [env.position]
        
        # 初始化可视化
        ax.clear()
        robot_paths = {0: path}
        goals = {0: env.goal}
        env.grid_map.render(fig=fig, ax=ax, robot_paths=robot_paths, goals=goals)
        ax.set_title(f"训练回合 {episode}, 步数: {num_steps}")
        plt.draw()
        plt.pause(0.5)
        
        while True:
            # 选择动作
            action = agent.select_action(state)
            
            # 执行动作
            next_state, reward, done, info = env.step(action)
            
            # 记录路径
            path.append(env.position)
            
            # 记录振荡
            if info.get('is_oscillating', False):
                oscillation_count += 1
            
            # 动态可视化
            if num_steps % 2 == 0 or done:  # 每两步更新一次可视化，减少刷新频率
                ax.clear()
                robot_paths = {0: path}
                goals = {0: env.goal}
                env.grid_map.render(fig=fig, ax=ax, robot_paths=robot_paths, goals=goals)
                status_text = f"训练回合: {episode}, 步数: {num_steps+1}"
                if info.get('is_oscillating', False):
                    status_text += " [检测到振荡!]"
                ax.set_title(status_text)
                plt.draw()
                plt.pause(0.1)
            
            # 存储经验
            agent.store_transition(state, action, reward, next_state, done)
            
            # 学习
            agent.learn()
            
            # 更新状态
            state = next_state
            episode_reward += reward
            num_steps += 1
            
            # 如果回合结束或者超过最大步数
            if done or num_steps >= 100:  # 设置较小的步数上限，避免无限循环
                break
        
        # 更新目标网络
        if episode % agent.target_update == 0:
            agent.update_target_network()
        
        # 更新探索率
        agent.update_epsilon()
        
        # 记录训练数据
        rewards_history.append(episode_reward)
        steps_history.append(num_steps)
        success = env.grid_map.is_done(env.position, env.goal)
        success_history.append(1 if success else 0)
        oscillation_counts.append(oscillation_count)
        
        # 打印训练信息
        success_rate = np.mean(success_history[-min(10, len(success_history)):]) * 100
        print(f"回合: {episode}/{episodes}, "
              f"奖励: {episode_reward:.2f}, "
              f"步数: {num_steps}, "
              f"振荡次数: {oscillation_count}, "
              f"成功: {'是' if success else '否'}, "
              f"成功率: {success_rate:.1f}%")
        
        # 在回合结束后暂停一下
        plt.pause(0.5)
    
    # 关闭交互模式
    plt.ioff()
    
    # 保存最终模型
    model_name = f"dqn_{'with' if use_fix else 'without'}_oscillation_fix.pth"
    os.makedirs("models", exist_ok=True)
    final_model_path = os.path.join("models", model_name)
    agent.save_model(final_model_path)
    print(f"最终模型已保存到 {final_model_path}")
    
    # 绘制训练结果
    plt.figure(figsize=(15, 10))
    
    plt.subplot(2, 2, 1)
    plt.plot(rewards_history)
    plt.title('回合奖励')
    plt.xlabel('回合')
    plt.ylabel('奖励')
    
    plt.subplot(2, 2, 2)
    plt.plot(steps_history)
    plt.title('回合步数')
    plt.xlabel('回合')
    plt.ylabel('步数')
    
    plt.subplot(2, 2, 3)
    window_size = min(10, len(success_history))
    if window_size > 0:
        success_moving_avg = [
            np.mean(success_history[max(0, i - window_size):i+1]) * 100 
            for i in range(len(success_history))
        ]
        plt.plot(success_moving_avg)
        plt.title('成功率 (移动平均)')
        plt.xlabel('回合')
        plt.ylabel('成功率 (%)')
        plt.ylim(0, 105)
    
    plt.subplot(2, 2, 4)
    plt.plot(oscillation_counts)
    plt.title('每回合振荡次数')
    plt.xlabel('回合')
    plt.ylabel('振荡次数')
    
    plt.tight_layout()
    plt.savefig(os.path.join("results", f"oscillation_fix_{'with' if use_fix else 'without'}.png"))
    plt.close()  # 不显示，只保存
    
    return agent, grid_map, {
        'rewards': rewards_history,
        'steps': steps_history,
        'success_rate': np.mean(success_history) * 100,
        'oscillation_counts': oscillation_counts
    }

def compare_oscillation_fix(episodes=50, seed=42):
    """比较有无振荡修复的效果差异"""
    # 训练不带修复的模型
    print("\n===== 训练不带振荡修复的模型 =====")
    _, _, results_without_fix = train_oscillation_fix(episodes=episodes, use_fix=False, seed=seed)
    
    # 训练带修复的模型
    print("\n===== 训练带振荡修复的模型 =====")
    _, _, results_with_fix = train_oscillation_fix(episodes=episodes, use_fix=True, seed=seed)
    
    # 比较结果
    plt.figure(figsize=(15, 10))
    
    # 比较奖励
    plt.subplot(2, 2, 1)
    plt.plot(results_without_fix['rewards'], label='不带修复')
    plt.plot(results_with_fix['rewards'], label='带修复')
    plt.title('回合奖励对比')
    plt.xlabel('回合')
    plt.ylabel('奖励')
    plt.legend()
    
    # 比较步数
    plt.subplot(2, 2, 2)
    plt.plot(results_without_fix['steps'], label='不带修复')
    plt.plot(results_with_fix['steps'], label='带修复')
    plt.title('回合步数对比')
    plt.xlabel('回合')
    plt.ylabel('步数')
    plt.legend()
    
    # 比较振荡次数
    plt.subplot(2, 2, 3)
    plt.plot(results_without_fix['oscillation_counts'], label='不带修复')
    plt.plot(results_with_fix['oscillation_counts'], label='带修复')
    plt.title('振荡次数对比')
    plt.xlabel('回合')
    plt.ylabel('振荡次数')
    plt.legend()
    
    # 比较平均指标
    plt.subplot(2, 2, 4)
    labels = ['平均奖励', '平均步数', '成功率(%)', '平均振荡次数']
    without_fix = [
        np.mean(results_without_fix['rewards']),
        np.mean(results_without_fix['steps']),
        results_without_fix['success_rate'],
        np.mean(results_without_fix['oscillation_counts'])
    ]
    with_fix = [
        np.mean(results_with_fix['rewards']),
        np.mean(results_with_fix['steps']),
        results_with_fix['success_rate'],
        np.mean(results_with_fix['oscillation_counts'])
    ]
    
    x = np.arange(len(labels))
    width = 0.35
    
    plt.bar(x - width/2, without_fix, width, label='不带修复')
    plt.bar(x + width/2, with_fix, width, label='带修复')
    plt.title('平均指标对比')
    plt.xticks(x, labels)
    plt.legend()
    
    plt.tight_layout()
    plt.savefig(os.path.join("results", "oscillation_fix_comparison.png"))
    plt.show()

def main():
    parser = argparse.ArgumentParser(description="测试振荡修复效果")
    parser.add_argument("--mode", type=str, default="compare", choices=["train", "compare"], 
                        help="运行模式: train(单独训练)或compare(比较两种方法)")
    parser.add_argument("--episodes", type=int, default=50, help="训练回合数")
    parser.add_argument("--use_fix", action="store_true", help="是否使用振荡修复")
    parser.add_argument("--seed", type=int, default=42, help="随机种子")
    
    args = parser.parse_args()
    
    if args.mode == "train":
        train_oscillation_fix(
            episodes=args.episodes,
            use_fix=args.use_fix,
            seed=args.seed
        )
    else:
        compare_oscillation_fix(
            episodes=args.episodes,
            seed=args.seed
        )

if __name__ == "__main__":
    main() 